#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time   : 2023-08-26 21:22
# @Author : hard work
# @File    : 练习13-parse_seq_num.py
# @Description :
"""
* merge fragment A and fragment B to C
* parse C.pcap pkt file
* compare seq num
"""
import os
import sys
from pyshark_util.pyshark_util import dec2addr, addr2dec
from parse_pkt_base import Parse_pcap_base


class Check_pkt_continuity(Parse_pcap_base):
    def __init__(self, file, file2, wireshark_file):
        super().__init__(file)
        self.__merge_file = r'D:\Tools\merge.pcap'
        self.__wireshark_file = wireshark_file

        self.__file_a = file
        self.__file_b = file2

        self.merge_frag_pcap()


    def merge_frag_pcap(self):
        try:
            if not os.path.exists(self.__file_a) and not os.path.exists(self.__file_b) \
                    and os.path.exists(self.__wireshark_file):
                raise "please check file1 and file2 or wireshark_file path"

            else:
                print("Info", "find file {}".format(self.__file_a))
                print("Info", "find file {}".format(self.__file_b))
                print("Info", "find file {}".format(self.__wireshark_file))

            if os.path.exists(self.__merge_file):
                print("Info", "file {} is existed , remove is first".format(self.__merge_file))
                os.remove(self.__merge_file)
            else:
                print("Info", "after merge path is {}".format(self.__merge_file))

            try:
                command = self.__wireshark_file + r"\mergecap.exe -w " + self.__merge_file + " " + self.__file_a \
                          + " " + self.__file_b
                """call system wireshark's mergecap.exe merge two files"""

                os.system(command)
                print("Info", "merge file  is done ...")
                print("Info", "right now parse and merge frag file, then take a few second ...")
                # return self.parse_pkt()
                self.parse_pcap_to_one_list(file_path=self.__merge_file, wireshark_file=self.__wireshark_file)
            except OSError as e:
                print("merge file error ", e)
            except RuntimeError as e:
                print("merge file error ", e)
        except FileNotFoundError as f:
            print("file not exits ", f)

    def compare_ip_seq_num(self):
        """判断ip的seq是否连续"""
        # pkt_lst = merge_and_parse_pkt()
        # merger_file = r"D:\cap\cb.pcap"
        # pkt_lst = parse_pkt(merger_file)
        ip_id_lst = []
        # res = []
        for pkt in self.PKT_LST:
            hex_2_int = int(pkt.ip.id, 16)
            # hex_2_int = int(pkt.IEEE8021CB.SEQ, 16)
            ip_id_lst.append(hex_2_int)
        self.compare_pkt_field_num(ip_id_lst)
        # cnt = 0
        # while cnt < len(ip_id_lst) - 1:
        #     if ip_id_lst[cnt + 1] - 1 == ip_id_lst[cnt]:
        #         self.mylog.NOTICE_LOG("第%s比较 %s 减去1 == %s " % (
        #             cnt, ip_id_lst[cnt + 1], ip_id_lst[cnt]))
        #     else:
        #         self.mylog.WARNNING_LOG("第%s比较 %s 减去1 != %s " % (
        #             cnt, ip_id_lst[cnt + 1], ip_id_lst[cnt]))
        #     cnt += 1
        # for i in range(len(ip_id)):
        #
        #     if not res:
        #         res.append([ip_id[i]])
        #     elif ip_id[i - 1] + 1 == ip_id[i]:
        #         res[-1].append(ip_id[i])
        #         print("Info: ", "第%s个pkt: %s " % ((count + 1), ip_id[i - 1]), "与第%s个pkt: %s 相差%s" %
        #               (count, ip_id[i], ip_id[i] - ip_id[i - 1]))
        #     else:
        #         res.append([ip_id[i]])
        #         print("\033[1;31m Error : ", "第%s个pkt: %s " % ((count + 1), ip_id[i - 1]), "与第%s个pkt: %s 相差%s \033[0m" %
        #               (count, ip_id[i], ip_id[i] - ip_id[i - 1]))
        #     count += 1
        # print("result -> ", res)

    def compare_eth_srcMac(self):
        # pkt_lst = merge_and_parse_pkt()
        mac_sa_lst = []
        res = []
        for pkt in self.PKT_LST:
            mac_sa = pkt.eth.src
            mac_sa = mac_sa.replace(":", "")
            hex_2_int = int(mac_sa, 16)
            mac_sa_lst.append(hex_2_int)
        self.compare_pkt_field_num(mac_sa_lst)
        # count = 0
        # for i in range(len(mac_sa_lst)):
        #
        #     if not res:
        #         res.append([mac_sa_lst[i]])
        #     elif mac_sa_lst[i - 1] + 1 == mac_sa_lst[i]:
        #         res[-1].append(mac_sa_lst[i])
        #         print("Info: ", "第%s个pkt: %s " % ((count + 1), mac_sa_lst[i - 1]), "与第%s个pkt: %s 相差%s" %
        #               (count, mac_sa_lst[i], mac_sa_lst[i] - mac_sa_lst[i - 1]))
        #     else:
        #         res.append([mac_sa_lst[i]])
        #         print("\033[1;31m Error : ", "第%s个pkt: %s " % ((count + 1), mac_sa_lst[i - 1]),
        #               "与第%s个pkt: %s 相差%s \033[0m" %
        #               (count, mac_sa_lst[i], mac_sa_lst[i] - mac_sa_lst[i - 1]))
        #     count += 1
        # print("result -> ", res)

    def compare_eth_dstMac(self):
        # pkt_lst = merge_and_parse_pkt()
        mac_da_lst = []
        res = []
        for pkt in self.PKT_LST:
            mac_sa = pkt.eth.dst
            mac_sa = mac_sa.replace(":", "")
            hex_2_int = int(mac_sa, 16)
            mac_da_lst.append(hex_2_int)
        self.compare_pkt_field_num(mac_da_lst)
        # count = 0
        # for i in range(len(mac_da_lst)):
        #
        #     if not res:
        #         res.append([mac_da_lst[i]])
        #     elif mac_da_lst[i - 1] + 1 == mac_da_lst[i]:
        #         res[-1].append(mac_da_lst[i])
        #         print("Info: ", "第%s个pkt: %s " % ((count + 1), mac_da_lst[i - 1]), "与第%s个pkt: %s 相差%s" %
        #               (count, mac_da_lst[i], mac_da_lst[i] - mac_da_lst[i - 1]))
        #     else:
        #         res.append([mac_da_lst[i]])
        #         print("\033[1;31m Error : ", "第%s个pkt: %s " % ((count + 1), mac_da_lst[i - 1]),
        #               "与第%s个pkt: %s 相差%s \033[0m" %
        #               (count, mac_da_lst[i], mac_da_lst[i] - mac_da_lst[i - 1]))
        #     count += 1
        # print("result -> ", res)

    def compare_ip_dst(self):
        # pkt_lst = merge_and_parse_pkt()
        ip_da_lst = []
        res = []
        for pkt in self.PKT_LST:
            ip_dst = pkt.ip.dst
            ip_da_lst.append(addr2dec(ip_dst))
        self.compare_pkt_field_num(ip_da_lst)
        # count = 0
        # for i in range(len(ip_da_lst)):
        #
        #     if not res:
        #         res.append([ip_da_lst[i]])
        #     elif ip_da_lst[i - 1] + 1 == ip_da_lst[i]:
        #         res[-1].append(ip_da_lst[i])
        #         print("Info: ", "第%s个pkt: %s " % ((count + 1), dec2addr(ip_da_lst[i - 1])), "与第%s个pkt: %s 相差 %s" %
        #               (count, dec2addr(ip_da_lst[i]), ip_da_lst[i] - ip_da_lst[i - 1]))
        #     else:
        #         res.append([ip_da_lst[i]])
        #         print("\033[1;31m Error : ", "第%s个pkt: %s " % ((count + 1), dec2addr(ip_da_lst[i - 1])),
        #               "与第%s个pkt: %s 相差 %s \033[0m" %
        #               (count, dec2addr(ip_da_lst[i]), ip_da_lst[i] - ip_da_lst[i - 1]))
        #     count += 1
        # print("result -> ", res)

    def compare_ip_src(self):
        # pkt_lst = merge_and_parse_pkt()
        ip_src_lst = []
        res = []
        for pkt in self.PKT_LST:
            ip_dst = pkt.ip.src
            ip_src_lst.append(addr2dec(ip_dst))
        self.compare_pkt_field_num(ip_src_lst)
        # count = 0
        # for i in range(len(ip_src_lst)):
        #
        #     if not res:
        #         res.append([ip_src_lst[i]])
        #     elif ip_src_lst[i - 1] + 1 == ip_src_lst[i]:
        #         res[-1].append(ip_src_lst[i])
        #         print("Info: ", "第%s个pkt: %s " % ((count + 1), dec2addr(ip_src_lst[i - 1])), "与第%s个pkt: %s 相差 %s" %
        #               (count, dec2addr(ip_src_lst[i]), ip_src_lst[i] - ip_src_lst[i - 1]))
        #     else:
        #         res.append([ip_src_lst[i]])
        #         print("\033[1;31m Error : ", "第%s个pkt: %s " % ((count + 1), dec2addr(ip_src_lst[i - 1])),
        #               "与第%s个pkt: %s 相差 %s \033[0m" %
        #               (count, dec2addr(ip_src_lst[i]), ip_src_lst[i] - ip_src_lst[i - 1]))
        #     count += 1
        # print("result -> ", res)

    def compare_tcp_seq_num(self):
        # pkt_lst = merge_and_parse_pkt()
        # merger_file = r"D:\Tools\tcp.pcap"
        # pkt_lst = parse_pkt(merger_file)
        tcp_id = []
        res = []
        for pkt in self.PKT_LST:
            hex_2_int = int(pkt.tcp.seq, 16)
            # hex_2_int = int(pkt.IEEE8021CB.SEQ, 16)
            tcp_id.append(hex_2_int)
        self.compare_pkt_field_num(tcp_id)

        # count = 0
        # for i in range(len(tcp_id)):
        #
        #     if not res:
        #         res.append([tcp_id[i]])
        #     elif tcp_id[i - 1] + 1 == tcp_id[i]:
        #         res[-1].append(tcp_id[i])
        #         print("Info: ", "第%s个pkt: %s " % ((count + 1), tcp_id[i - 1]), "与第%s个pkt: %s 相差%s" %
        #               (count, tcp_id[i], tcp_id[i] - tcp_id[i - 1]))
        #     else:
        #         res.append([tcp_id[i]])
        #         print("\033[1;31m Error : ", "第%s个pkt: %s " % ((count + 1), tcp_id[i - 1]),
        #               "与第%s个pkt: %s 相差%s \033[0m" %
        #               (count, tcp_id[i], tcp_id[i] - tcp_id[i - 1]))
        #     count += 1
        # print("result -> ", res)

    def compare_stc_sig(self):
        # pkt_lst = merge_and_parse_pkt()
        # merger_file = r"D:\Tools\3.pcap"
        # pkt_lst = parse_pkt(merger_file)
        stc_sig_lst = []
        res = []
        for pkt in self.PKT_LST:
            # print(dir(pkt.STCSIG))
            # print(pkt.STCSIG.seqnum)
            hex_2_int = int(pkt.STCSIG.seqnum, 16)
            stc_sig_lst.append(hex_2_int)
        self.compare_pkt_field_num(stc_sig_lst)

        # count = 0
        # for i in range(len(stc_sig_lst)):
        #
        #     if not res:
        #         res.append([stc_sig_lst[i]])
        #     elif stc_sig_lst[i - 1] + 1 == stc_sig_lst[i]:
        #         res[-1].append(stc_sig_lst[i])
        #         print("Info: ", "第%s个pkt: %s " % ((count + 1), stc_sig_lst[i - 1]), "与第%s个pkt: %s 相差%s" %
        #               (count, stc_sig_lst[i], stc_sig_lst[i] - stc_sig_lst[i - 1]))
        #     else:
        #         res.append([stc_sig_lst[i]])
        #         print("\033[1;31m Error : ", "第%s个pkt: %s " % ((count + 1), stc_sig_lst[i - 1]),
        #               "与第%s个pkt: %s 相差%s \033[0m" %
        #               (count, stc_sig_lst[i], stc_sig_lst[i] - stc_sig_lst[i - 1]))
        #     count += 1
        # print("result -> ", res)


def main():
    print("welcome use this tool, it's will help to you")
    print(' merge fragment A and fragment B to C ')
    print("\033[1;31m Warning:You need to modify the file path to be real configuration\033[0m")
    print("0: compare src_mac   1: compare dst_mac   2: compare src_ip   3: compare dst_ip")
    print("4: compare ip_id   5: compare tcp_seq   6: compare stc_sig")
    print("7: exit")

    USER_INPUT = int(input("Please enter the number selection function : "))
    try:
        if USER_INPUT >= 7 or USER_INPUT < 0:
            print("retry")
            sys.exit(1)
        else:
            # file1 = r"D:\Tools\3.pcap"
            file1 = input("请输入文件路径>>>")
            # file2 = r"D:\Tools\4.pcap"
            file2 = input("请输入文件路径>>>")
            # merger_file = r"D:\Tools\merge.pcap"
            # wireshark_file = r"D:\\Tools\\Wireshark"
            wireshark_file = input("请输入wireshark文件路径>>>")
            if os.path.isfile(file1.strip()) or os.path.isfile(file2.strip()) or os.path.isfile(wireshark_file.strip()):
                _ckeck = Check_pkt_continuity(file1, file2, wireshark_file)
                if USER_INPUT == 0:
                    _ckeck.compare_eth_srcMac()
                elif USER_INPUT == 1:
                    _ckeck.compare_eth_dstMac()
                elif USER_INPUT == 2:
                    _ckeck.compare_ip_src()
                elif USER_INPUT == 3:
                    _ckeck.compare_ip_dst()
                elif USER_INPUT == 4:
                    print("may be need to disable ddos setting")
                    _ckeck.compare_ip_seq_num()
                elif USER_INPUT == 5:
                    print("may be need to disable ddos setting")
                    _ckeck.compare_tcp_seq_num()
                elif USER_INPUT == 6:
                    _ckeck.compare_stc_sig()

            else:
                raise 'file not found'
    except Exception as e:
        print(e)
        sys.exit(1)


if __name__ == '__main__':
    main()
